"""符号矩阵模块
包含符号计算支持、矩阵符号导数等功能
"""
import numpy as np
import sympy as sp
from typing import List, Dict, Union, Tuple, Optional
from sympy import symbols, Matrix, Symbol, simplify, expand, factor
from sympy import eye, zeros, ones, diag
from sympy.matrices import MatrixSymbol

class SymbolicMatrix:
    @staticmethod
    def create_symbolic_matrix(rows: int, cols: int, 
                            prefix: str = 'm') -> sp.Matrix:
        """
        创建符号矩阵
        :param rows: 行数
        :param cols: 列数
        :param prefix: 符号前缀
        :return: SymPy符号矩阵
        """
        symbols = sp.symbols(f'{prefix}_:{rows*cols}')
        return sp.Matrix(symbols).reshape(rows, cols)
    
    @staticmethod
    def create_matrix_symbol(name: str, rows: int, cols: int) -> MatrixSymbol:
        """
        创建矩阵符号
        :param name: 矩阵名称
        :param rows: 行数
        :param cols: 列数
        :return: 矩阵符号
        """
        return MatrixSymbol(name, rows, cols)
    
    @staticmethod
    def create_special_matrices(n: int) -> Dict[str, sp.Matrix]:
        """
        创建特殊符号矩阵
        :param n: 矩阵维度
        :return: 特殊矩阵字典
        """
        return {
            'identity': eye(n),
            'zeros': zeros(n, n),
            'ones': ones(n, n),
            'diagonal': diag(*symbols(f'd_:{n}')),
            'symmetric': SymbolicMatrix.create_symmetric_matrix(n),
            'antisymmetric': SymbolicMatrix.create_antisymmetric_matrix(n)
        }
    
    @staticmethod
    def create_identity_matrix(n: int) -> sp.Matrix:
        """
        创建单位矩阵
        :param n: 矩阵维度
        :return: 单位矩阵
        """
        return eye(n)
    
    @staticmethod
    def create_zeros_matrix(rows: int, cols: int) -> sp.Matrix:
        """
        创建零矩阵
        :param rows: 行数
        :param cols: 列数
        :return: 零矩阵
        """
        return zeros(rows, cols)
    
    @staticmethod
    def create_ones_matrix(rows: int, cols: int) -> sp.Matrix:
        """
        创建全1矩阵
        :param rows: 行数
        :param cols: 列数
        :return: 全1矩阵
        """
        return ones(rows, cols)
    
    @staticmethod
    def create_symmetric_matrix(n: int, prefix: str = 's') -> sp.Matrix:
        """
        创建对称符号矩阵
        :param n: 矩阵维度
        :param prefix: 符号前缀
        :return: 对称矩阵
        """
        # 创建上三角元素
        elements = []
        for i in range(n):
            for j in range(i, n):
                elements.append(symbols(f'{prefix}_{i}_{j}'))
        
        # 构造对称矩阵
        matrix = zeros(n, n)
        idx = 0
        for i in range(n):
            for j in range(i, n):
                matrix[i, j] = elements[idx]
                if i != j:
                    matrix[j, i] = elements[idx]
                idx += 1
        
        return matrix
    
    @staticmethod
    def create_antisymmetric_matrix(n: int, prefix: str = 'a') -> sp.Matrix:
        """
        创建反对称符号矩阵
        :param n: 矩阵维度
        :param prefix: 符号前缀
        :return: 反对称矩阵
        """
        # 创建上三角元素（不包括对角线）
        elements = []
        for i in range(n):
            for j in range(i + 1, n):
                elements.append(symbols(f'{prefix}_{i}_{j}'))
        
        # 构造反对称矩阵
        matrix = zeros(n, n)
        idx = 0
        for i in range(n):
            for j in range(i + 1, n):
                matrix[i, j] = elements[idx]
                matrix[j, i] = -elements[idx]
                idx += 1
        
        return matrix
    
    @staticmethod
    def matrix_derivative(matrix: sp.Matrix, 
                        var: sp.Symbol) -> sp.Matrix:
        """
        计算矩阵对符号变量的导数
        :param matrix: 输入矩阵
        :param var: 求导变量
        :return: 导数矩阵
        """
        return matrix.diff(var)
    
    @staticmethod
    def matrix_jacobian(matrix: sp.Matrix, variables: List[sp.Symbol]) -> sp.Matrix:
        """
        计算矩阵的雅可比矩阵
        :param matrix: 输入矩阵（向量）
        :param variables: 变量列表
        :return: 雅可比矩阵
        """
        if matrix.cols != 1:
            raise ValueError("Input matrix must be a column vector")
        
        jacobian = zeros(matrix.rows, len(variables))
        for i in range(matrix.rows):
            for j, var in enumerate(variables):
                jacobian[i, j] = matrix[i].diff(var)
        
        return jacobian
    
    @staticmethod
    def matrix_hessian(expr: sp.Expr, variables: List[sp.Symbol]) -> sp.Matrix:
        """
        计算标量函数的Hessian矩阵
        :param expr: 标量表达式
        :param variables: 变量列表
        :return: Hessian矩阵
        """
        n = len(variables)
        hessian = zeros(n, n)
        
        for i in range(n):
            for j in range(n):
                hessian[i, j] = expr.diff(variables[i], variables[j])
        
        return hessian
    
    @staticmethod
    def symbolic_eigenvalues(matrix: sp.Matrix) -> List[sp.Expr]:
        """
        计算符号矩阵的特征值
        :param matrix: 输入矩阵
        :return: 特征值列表
        """
        if matrix.rows != matrix.cols:
            raise ValueError("Matrix must be square")
        
        return list(matrix.eigenvals().keys())
    
    @staticmethod
    def symbolic_eigenvects(matrix: sp.Matrix) -> List[Tuple[sp.Expr, int, List[sp.Matrix]]]:
        """
        计算符号矩阵的特征向量
        :param matrix: 输入矩阵
        :return: (特征值, 重数, 特征向量列表)的列表
        """
        if matrix.rows != matrix.cols:
            raise ValueError("Matrix must be square")
        
        return matrix.eigenvects()
    
    @staticmethod
    def characteristic_polynomial(matrix: sp.Matrix, var: sp.Symbol = None) -> sp.Poly:
        """
        计算特征多项式
        :param matrix: 输入矩阵
        :param var: 多项式变量
        :return: 特征多项式
        """
        if matrix.rows != matrix.cols:
            raise ValueError("Matrix must be square")
        
        if var is None:
            var = symbols('lambda')
        
        return matrix.charpoly(var)
    
    @staticmethod
    def minimal_polynomial(matrix: sp.Matrix, var: sp.Symbol = None) -> sp.Poly:
        """
        计算最小多项式
        :param matrix: 输入矩阵
        :param var: 多项式变量
        :return: 最小多项式
        """
        if matrix.rows != matrix.cols:
            raise ValueError("Matrix must be square")
        
        if var is None:
            var = symbols('lambda')
        
        return matrix.minpoly(var)
    
    @staticmethod
    def matrix_function(matrix: sp.Matrix, func_name: str) -> sp.Matrix:
        """
        计算矩阵函数
        :param matrix: 输入矩阵
        :param func_name: 函数名称 ('exp', 'log', 'sin', 'cos', 'sqrt')
        :return: 矩阵函数结果
        """
        if matrix.rows != matrix.cols:
            raise ValueError("Matrix must be square")
        
        if func_name == 'exp':
            return matrix.exp()
        elif func_name == 'log':
            # 对于符号矩阵，对数函数需要特殊处理
            eigenvals = matrix.eigenvals()
            if all(val.is_positive for val in eigenvals.keys()):
                # 如果所有特征值都是正数，可以计算对数
                P, D = matrix.diagonalize()
                log_D = diag(*[sp.log(d) for d in D.diagonal()])
                return P * log_D * P.inv()
            else:
                raise ValueError("Matrix logarithm requires positive eigenvalues")
        elif func_name == 'sqrt':
            # 矩阵平方根
            eigenvals = matrix.eigenvals()
            if all(val.is_nonnegative for val in eigenvals.keys()):
                P, D = matrix.diagonalize()
                sqrt_D = diag(*[sp.sqrt(d) for d in D.diagonal()])
                return P * sqrt_D * P.inv()
            else:
                raise ValueError("Matrix square root requires non-negative eigenvalues")
        else:
            raise ValueError(f"Unsupported function: {func_name}")
    
    @staticmethod
    def matrix_power(matrix: sp.Matrix, power: Union[int, sp.Expr]) -> sp.Matrix:
        """
        计算矩阵幂
        :param matrix: 输入矩阵
        :param power: 幂次
        :return: 矩阵幂
        """
        if matrix.rows != matrix.cols:
            raise ValueError("Matrix must be square")
        
        return matrix ** power
    
    @staticmethod
    def simplify_matrix(matrix: sp.Matrix) -> sp.Matrix:
        """
        简化符号矩阵
        :param matrix: 输入矩阵
        :return: 简化后的矩阵
        """
        return matrix.applyfunc(simplify)
    
    @staticmethod
    def expand_matrix(matrix: sp.Matrix) -> sp.Matrix:
        """
        展开符号矩阵
        :param matrix: 输入矩阵
        :return: 展开后的矩阵
        """
        return matrix.applyfunc(expand)
    
    @staticmethod
    def factor_matrix(matrix: sp.Matrix) -> sp.Matrix:
        """
        因式分解符号矩阵
        :param matrix: 输入矩阵
        :return: 因式分解后的矩阵
        """
        return matrix.applyfunc(factor)
    
    @staticmethod
    def substitute_matrix(matrix: sp.Matrix, substitutions: Dict[sp.Symbol, Union[sp.Expr, float]]) -> sp.Matrix:
        """
        符号矩阵变量替换
        :param matrix: 输入矩阵
        :param substitutions: 替换字典
        :return: 替换后的矩阵
        """
        return matrix.subs(substitutions)
    
    @staticmethod
    def matrix_trace_symbolic(matrix: sp.Matrix) -> sp.Expr:
        """
        计算符号矩阵的迹
        :param matrix: 输入矩阵
        :return: 矩阵的迹
        """
        if matrix.rows != matrix.cols:
            raise ValueError("Matrix must be square")
        
        return matrix.trace()
    
    @staticmethod
    def matrix_determinant_symbolic(matrix: sp.Matrix) -> sp.Expr:
        """
        计算符号矩阵的行列式
        :param matrix: 输入矩阵
        :return: 矩阵的行列式
        """
        if matrix.rows != matrix.cols:
            raise ValueError("Matrix must be square")
        
        return matrix.det()
    
    @staticmethod
    def matrix_inverse_symbolic(matrix: sp.Matrix) -> sp.Matrix:
        """
        计算符号矩阵的逆
        :param matrix: 输入矩阵
        :return: 矩阵的逆
        """
        if matrix.rows != matrix.cols:
            raise ValueError("Matrix must be square")
        
        return matrix.inv()
    
    @staticmethod
    def solve_matrix_equation(A: sp.Matrix, B: sp.Matrix) -> sp.Matrix:
        """
        求解矩阵方程 AX = B
        :param A: 系数矩阵
        :param B: 常数矩阵
        :return: 解矩阵X
        """
        if A.rows != A.cols:
            raise ValueError("Coefficient matrix A must be square")
        if A.rows != B.rows:
            raise ValueError("Matrices A and B must have the same number of rows")
        
        return A.inv() * B